#!/usr/bin/env python
"""

    index_fasta.py < sequence.fa > sequence.idx
    [-i, --input_file -]
    [-w, --whitespace_delimited]
    [--log_file PATH]
    [--quiet]
    
"""
#
# These contig names are equivalent but referred to differently in Ensembl / ucsc
contig_aliases =    {
                        "chrM"           : "chrMT"        ,
                        "mitochondria"   : "chrMT"        ,
                        "chloroplast"    : "chrC"        ,
                        "chr6_cox_hap1"  : "chr6_COX"     ,
                        "chr6_qbl_hap2"  : "chr6_QBL"     ,
                        "chr5_h2_hap1"   : "chr5_H2,"     ,
                        "chr22_h2_hap1"  : "chr22_H2"     ,
                        "chrMT"          : "chrM"         ,
                        "chr6_COX"       : "chr6_cox_hap1",
                        "chr6_QBL"       : "chr6_qbl_hap2",
                        "chr5_H2,"       : "chr5_h2_hap1" ,
                        "chr22_H2"       : "chr22_h2_hap1",
                    }


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   options        


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

import optparse
import sys, os
import os.path
import StringIO

# add self to search path for testing
if __name__ == '__main__':
    exe_path = os.path.split(os.path.abspath(sys.argv[0]))[0]
    python_modules_path = os.path.abspath(os.path.join(exe_path,"..", "python_modules"))
    sys.path.append(python_modules_path)
    module_name = os.path.split(sys.argv[0])[1]
    module_name = os.path.splitext(module_name)[0];
else:
    module_name = __name__



if __name__ == '__main__':
    parser = optparse.OptionParser(version="%prog 1.0", usage = "\n\n    %progs [options]")
    parser.add_option("-i", "--input_file", dest="input_file",
                      metavar="FILE", 
                      type="string",
                      default = "-",
                      help="Name and path of sequence file in FASTA format. "
                            "Defaults to reading from STDIN.")

    parser.add_option("-o", "--output_file", dest="output_file",
                      metavar="FILE", 
                      default = "-",
                      type="string",
                      help="name and path of index file. Defaults to stdout.")
    
    parser.add_option("-w", "--whitespace_delimited", dest = "whitespace_delimited",
                      action="store_true", default=False,
                      help="Contig name is accession line only up to first whitespace. "
                            "Otherwise, entire accession line is treated as the contig name. ")
    
    
    
    #
    #   general options: verbosity / logging
    # 
    parser.add_option("-v", "--verbose", dest = "verbose",
                      action="store_true", default=False,
                      help="Print progress to STDERR")
    parser.add_option("-L", "--log_file", dest="log_file",
                      metavar="FILE", 
                      type="string",
                      help="name and path of log file")
    parser.add_option("--skip_parameter_logging", dest="skip_parameter_logging",
                      action="store_true", default=False,
                      help="Do not print program parameters to log.")
    parser.add_option("--debug", dest="debug",
                        action="store_true", default=False,
                        help="Set default program parameters in debugging mode.")
    
    
    
    
    
    # get help string
    f =StringIO.StringIO()
    parser.print_help(f)
    helpstr = f.getvalue()
    (options, remaining_args) = parser.parse_args()
    
    if options.debug:
        options.verbose = True
        options.log_file = os.path.join("index_fasta.log")
    
    # mandatory options
    from options import check_mandatory_options
    mandatory_options = []
    check_mandatory_options (options, mandatory_options, helpstr)

#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   imports        


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

import StringIO
import re
import operator

from collections import defaultdict



#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   Functions        


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888
def index_fasta (seq_file, verbose = False, whitespace_delimited = True):
    genomic_sequence_coordinates = list()

    has_accession = False

    dot_count = 0
    new_acc_beg = 0
    new_acc_len = 0
    curr_file_pos = 0
    if verbose:
        sys.stderr.write("\t[");
    for line_count, line in enumerate(seq_file):

        if verbose and line_count % 5000 == 0:
            dot_count += 1
            if dot_count % 50 == 0:
                sys.stderr.write("\n")
            sys.stderr.write(".")

        curr_file_pos += len(line)
        if line[0] != '>':
            continue

        # only save things on accession line
        #curr_file_pos = seq_file.tell()

        # minus 1 for '>', 1 for newline
        new_acc_len = len(line) - 2;
        new_acc_beg = curr_file_pos - new_acc_len - 2 

        #
        #   Save data for previous accession
        # 
        if has_accession and acc_name != None:
            # seq_begin is the position of the first character after the '>'
            # 1 for '>', 1 for newline
            seq_len = new_acc_beg - seq_begin - 1;
            if seq_len > 0:
                if whitespace_delimited:
                    contig_name = acc_name[1:].split()[0]
                else:
                    contig_name = acc_name[1:]
                yield([contig_name, 
                        acc_begin, acc_len, 
                        seq_begin, seq_len])
                if contig_name in contig_aliases:
                    yield([contig_aliases[contig_name], 
                            acc_begin, acc_len, 
                            seq_begin, seq_len])
                    

        acc_begin, acc_len = new_acc_beg, new_acc_len

        # mark to say that we have encountered an accession line
        has_accession = True;

        #
        #   Save data for current accession
        # 

        # seq_begin is the position of the first character after the '>'
        seq_begin = curr_file_pos

        acc_name = line.strip()

    #
    #   Save data for last accession
    # 
    if has_accession and acc_name != None:
        curr_file_pos = curr_file_pos

        # seq_begin is the position of the first character after the '>'
        # 1 for '>', 1 for newline
        seq_len = curr_file_pos - seq_begin - 1;
        if seq_len > 0:
            if whitespace_delimited:
                contig_name = acc_name[1:].split()[0]
            else:
                contig_name = acc_name[1:]
            yield([contig_name, 
                    acc_begin, acc_len, 
                    seq_begin, seq_len])
            if contig_name in contig_aliases:
                yield([contig_aliases[contig_name], 
                        acc_begin, acc_len, 
                        seq_begin, seq_len])

    if verbose:
        sys.stderr.write("]\n");
#_________________________________________________________________________________________

#   write_fasta_index_to_file

#_________________________________________________________________________________________
def write_fasta_index_to_file (input_seq_file_name, output_index_file_name = "-", 
                                logger = None, verbose = False, whitespace_delimited = True):
      """
      writes index of FASTA format input to specified output

      Use "-" for input and output files to indicate STDIN/STDOUT    

      Writes the number of indexed sequences to logger

      """

      if input_seq_file_name == "-":
          logger.info("Reading from stdin")
          input_seq_file = sys.stdin
      else:
          logger.info("Input file = [%s]" % input_seq_file_name)
          input_seq_file = open(input_seq_file_name)

      if output_index_file_name == "-":
          logger.info("Writing to stdout")
          output_index_file = sys.stdout
      else:
          logger.info("Output file = [%s]" % output_index_file_name)
          output_index_file = open(output_index_file_name, "w")

      cnt = -1
      for cnt, sequence in enumerate(index_fasta (input_seq_file, verbose, whitespace_delimited)):
          output_index_file.write("\t".join(str(sequence[i]) for i in (0,1,3,4)) + "\n")

      if logger:
          message = "%d sequence%s indexed in fasta file" % (cnt + 1, 
                                                            "s" if cnt else "")
          if input_seq_file_name != "-":
              message += " [%s]" % input_seq_file_name
          logger.info(message)
    
    


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   Logger


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

import unittest
import logging
from lg_program_logging import setup_std_logging

if __name__ == "__main__":
    from options import get_option_strings
    logger = logging.getLogger(module_name)
    setup_std_logging(logger, options.log_file, options.verbose)

    if not options.skip_parameter_logging:
        programme_name = os.path.split(sys.argv[0])[1]
        logger.info("%s %s" % (programme_name, " ".join(get_option_strings(parser, options))))



#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   Main logic


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888



if __name__ == '__main__':
    if options.debug:
        sys.argv = sys.argv[:1]
        unittest.main()
        sys.exit(-1)
        
        
        
    write_fasta_index_to_file (options.input_file, options.output_file, logger, options.verbose,
                                options.whitespace_delimited)
        

